import queue
import logging
import threading
import av
import cv2
import numpy as np
import fractions


class Streamer:
    def __init__(self):
        self.run = False
        self.__log_init()

    def __log_init(self):
        self.logger = logging.getLogger("Streamer")

    def open(self, url):
        self.logger.trace("open.enter")
        self.container = av.open(url, mode="w", format="flv")
        self.stream = self.container.add_stream("libx264", rate=30)
        self.stream.width = 640
        self.stream.height = 480
        self.stream.pix_fmt = "yuv420p"
        self.stream.options = {
            "preset": "ultrafast",
            "tune": "zerolatency",
            "g": "10",
            "keyint_min": "1",
        }
        self.logger.info(f"rtmp: {url}")
        self.logger.trace("open.done")

    def close(self):
        self.logger.trace("close.enter")
        self.container.close()
        self.logger.trace("close.done")

    def thread(self):
        self.logger.trace("thread.enter")
        number = 0

        while self.run:
            try:
                item, _ = self.queue.get(timeout=1)
                frame = item["frame"]
                pts = item["timestamp"]
            except queue.Empty:
                self.logger.warning("queue get timeout!")
                continue
            except Exception as e:
                self.logger.error("queue get error: %s", str(e))
                continue

            number += 1
            rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            video_frame = av.VideoFrame.from_ndarray(rgb_frame, format="rgb24")
            video_frame.pts = pts
            video_frame.time_base = fractions.Fraction(1, 1000)

            self.logger.info(f"Rtmp   frame: {number}, timestamp: {pts}")

            try:
                for packet in self.stream.encode(video_frame):
                    self.container.mux(packet)
                self.logger.trace("frame send: %d", number)
            except Exception as e:
                self.logger.error("pyav mux exception: %s", str(e))
                break

        # flush 编码器（输出缓存帧）
        for packet in self.stream.encode():
            self.container.mux(packet)

        self.logger.trace("thread.done")

    def start(self):
        self.logger.trace("start.enter")
        self.run = True
        self.handle = threading.Thread(target=self.thread)
        self.handle.start()
        self.logger.trace("start.done")
        return True

    def stop(self):
        self.logger.trace("stop.enter")
        self.run = False
        self.handle.join()
        self.logger.trace("stop.done")

    def input_queue_set(self, queue):
        self.logger.trace("input_queue_set.enter")
        if self.run:
            self.logger.error("streamer is started, please stop the streamer first!")
            return False
        self.queue = queue
        self.logger.trace("input_queue_set.done")
        return True
